import os
import pandas as pd
from tabula import read_pdf
from openpyxl import load_workbook
import fitz  # PyMuPDF
import re
import math
from tkinter import filedialog
import tkinter as tk


# 从PDF中提取表格数据
def extract_table_data(pdf_path):
    dfs = read_pdf(pdf_path, pages="all", multiple_tables=True)
    df = pd.concat(dfs, ignore_index=True)
    filtered_df = df[(df['样品编号'].astype(str).str.startswith('2')) | (df['样品编号'].astype(str).str.contains('质控'))]
    filtered_df = filtered_df[['样品编号', 'Abs']]
    filtered_df.columns = [None, None]
    filtered_df.loc[-1] = ["以下空白", ""]
    filtered_df.index = range(len(filtered_df))
    return filtered_df


# 从PDF第2页提取Abs值
def extract_abs_values(pdf_path):
    page_2_data = read_pdf(pdf_path, pages=2, lattice=False, stream=True)
    combined_df = pd.DataFrame()
    for data_list in page_2_data:
        temp_df = pd.DataFrame(data_list)
        combined_df = pd.concat([combined_df, temp_df], ignore_index=True)
    column_names = combined_df.columns.tolist()
    new_data = {col: [col] + list(combined_df[col]) for col in column_names}
    new_df = pd.DataFrame(new_data)
    abs_column = new_df.iloc[:, 1]
    abs_values = abs_column.dropna().values
    return abs_values


# 从PDF中提取特定数值
def extract_numbers(pdf_path):
    doc = fitz.open(pdf_path)
    full_text = ""
    for page in doc:
        full_text += page.get_text("text")

    k1_match = re.search(r'K1\s*=\s*(-?\d*\.?\d+)', full_text)
    k0_match = re.search(r'K0\s*=\s*(-?\d*\.?\d+)', full_text)
    correlation_matches = re.findall(r'相关性：\s*([\d.,-]+)', full_text)

    k1_number = float(k1_match.group(1)) if k1_match.group(1) else None
    k0_number = float(k0_match.group(1)) if k0_match.group(1) else None
    correlation_numbers = [float(num.replace(',', '')) for num in correlation_matches[:3]]

    return k1_number, k0_number, correlation_numbers


# 只舍不入的四位有效数字的自定义舍入函数
def round_to_four_significant_digits(number):
    if not isinstance(number, (int, float)):
        raise ValueError("Input must be an integer or float.")

    # 获取数量级
    order_of_magnitude = int(math.floor(math.log10(abs(number)))) if number != 0 else 0
    precision = -(order_of_magnitude - 3)

    # 计算舍入因子
    factor = 10 ** precision

    # 向下取整
    rounded_number = math.floor(number * factor) / factor

    return rounded_number


# 将提取的数据写入Excel
def write_to_excel(filtered_df, abs_values, k1, k0, correlation, excel_path, sheet_names):
    wb = load_workbook(excel_path)

    # 写入到指定的工作表
    for sheet_name in sheet_names:
        ws = wb[sheet_name]
        if sheet_name.endswith('2') or sheet_name[len(sheet_name)-2]=="2":
            for i, row in enumerate(filtered_df.itertuples(), start=5):
                ws.cell(row=i, column=2, value=row[1])  # 写入样品编号
                ws.cell(row=i, column=5, value=row[2])  # 写入Abs值
        else:
            for j, abs_val in enumerate(abs_values, start=1):
                col_range_start = chr(64 + ((j - 1) * 2 + 3))  # 计算起始列字母
                col_range_end = chr(64 + ((j - 1) * 2 + 4))  # 计算结束列字母
                if j < 7:
                    ws.merge_cells(start_row=17, start_column=(j - 1) * 2 + 3, end_row=17, end_column=(j - 1) * 2 + 4)
                    ws.cell(row=17, column=(j - 1) * 2 + 3, value=abs_val)
                if len(abs_values) > 6:
                    ws.cell(row=17, column=15, value=abs_values[6])
                else:
                    ws.cell(row=17, column=15, value="/")

            # 写入K1值到G18:H18合并单元格
            ws.merge_cells(start_row=18, start_column=7, end_row=18, end_column=8)
            ws.cell(row=18, column=7).value = k1

            # 写入K0值到K18:L18合并单元格
            ws.merge_cells(start_row=18, start_column=11, end_row=18, end_column=12)
            ws.cell(row=18, column=11).value = k0

            # 写入相关系数到O18
            ws.cell(row=18, column=15).value = correlation

    wb.save(excel_path)


# 主程序
def run():
    # 使用filedialog选择PDF文件夹和目标Excel文件
    root = tk.Tk()
    root.withdraw()  # 隐藏主窗口
    pdf_folder = filedialog.askdirectory(title="选择PDF文件夹")
    excel_path = filedialog.askopenfilename(title="选择目标Excel文件", filetypes=[("Excel files", "*.xlsx")])

    if not pdf_folder or not excel_path:
        print("未选择文件夹或Excel文件，程序退出。")
        return

    # 定义PDF文件与工作表的对应关系
    pdf_sheet_map = {
        "铬": ["铬（原吸）", "铬（原吸）2"],
        "铁": ["铁", "铁2"],
        "铜": ["铜", "铜2"],
        "锌": ["锌", "锌2"],
        "铍":["铍","铍2"],
        "钠":["钠","钠2"],
        "锰": ["锰", "锰2"],
        "锰5750": ["锰5750", "锰5750（2）"],
        "铝5750": ["铝5750", "铝5750（2）"],
        "铁5750": ["铁5750", "铁5750（2）"],
        "钡": ["钡", "钡2"],
        "钙": ["钙", "钙2"],#曲线点减少一个
        "镁": ["镁", "镁2"],#曲线点减少一个
        "铅": ["铅", "铅2"],
        "镉": ["镉", "镉2"],
        "镉（废水）": ["镉（废水）", "镉（废水）2"],
        "铅（废水）": ["铅（废水）", "铅（废水）2"],
        "镍（废水）": ["镍（废水）", "镍（废水）2"],
        "镍（地下水）": ["镍（地下水）", "镍（地下水）2"],#曲线点增加一个
    }

    # 处理每个PDF文件
    for element, sheet_names in pdf_sheet_map.items():
        pdf_file = os.path.join(pdf_folder, f"{element}.pdf")
        if os.path.exists(pdf_file):
            # 提取表格数据
            filtered_df = extract_table_data(pdf_file)

            # 提取Abs值
            abs_values = extract_abs_values(pdf_file)

            # 提取特定数值
            k1, k0, correlations = extract_numbers(pdf_file)

            # 转换字符串为浮点数并应用自定义舍入
            first_correlation_float = float(correlations[0]) if correlations else None
            rounded_first_correlation = round_to_four_significant_digits(first_correlation_float) if first_correlation_float else None

            # 写入Excel
            write_to_excel(filtered_df, abs_values, k1, k0, rounded_first_correlation, excel_path, sheet_names)

    print("数据已成功写入Excel文件")


if __name__ == "__main__":
    run()